{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import paddle\n",
    "\n",
    "from paddle_harmonics.quadrature import legendre_gauss_weights, clenshaw_curtiss_weights\n",
    "from paddle_harmonics.legendre import legpoly, clm\n",
    "from paddle_harmonics import RealSHT, InverseRealSHT\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from plotting import plot_sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlat = 32\n",
    "nlon = 2*nlat\n",
    "grid = \"legendre-gauss\"\n",
    "\n",
    "# for quadrature and plotting\n",
    "if grid == \"legendre-gauss\":\n",
    "    lmax = mmax = nlat\n",
    "    xq, wq = legendre_gauss_weights(nlat)\n",
    "elif grid ==\"equiangular\":\n",
    "    lmax = mmax = nlat//2-1\n",
    "    xq, wq = clenshaw_curtiss_weights(nlat)\n",
    "\n",
    "sht = RealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid)\n",
    "isht = InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid)\n",
    "\n",
    "lat = np.arccos(xq)\n",
    "omega = math.pi * paddle.to_tensor(wq).to(paddle.float32) / nlat\n",
    "omega = omega.reshape([-1, 1]).tile([1, nlon])\n",
    "\n",
    "omega.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forming the Vandermonde matrix\n",
    "nmodes = int(lmax * (lmax+1) / 2)\n",
    "e = paddle.zeros([nmodes, lmax, mmax], dtype=paddle.complex64)\n",
    "\n",
    "midx = lambda l,m : l*(l+1)//2 + m\n",
    "\n",
    "for l in range(lmax):\n",
    "    for m in range(l+1):\n",
    "        e[midx(l,m), l, m] = 1.0\n",
    "\n",
    "vdm = isht(e) # nmodes x nlat x nlon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let us visualize the some entries in the Vandermonde matrix. In principle, it should contain the Spherical Harmonics in ascending order:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_lmax = 6\n",
    "\n",
    "fig = plt.figure(layout='constrained', figsize=(12, 8))\n",
    "subfigs = fig.subfigures(plt_lmax, plt_lmax)\n",
    "\n",
    "for l in range(plt_lmax):\n",
    "    for m in range(l+1):\n",
    "        plot_sphere(vdm[midx(l,m)], fig=subfigs[l, m], title=f\"l={l}, m={m}\", lat=(lat - np.pi/2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "now, let us compute the gramian matrix $S^T \\mathop{\\mathrm{diag}}(\\omega) S$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gramian = paddle.einsum(\"iqr,jqr,qr->ij\", vdm, vdm, omega)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.pcolormesh(xlg, np.arange(0, nmodes), , cmap=\"plasma\")\n",
    "# plt.xlabel(\"x\")\n",
    "# plt.ylabel(\"l,m mode\")\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(gramian, interpolation=\"nearest\", cmap=\"plasma\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(paddle.max(gramian))\n",
    "print(paddle.argmax(gramian))\n",
    "print(paddle.min(gramian))\n",
    "print(paddle.argmin(gramian))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linalg.cond(gramian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.diag(gramian))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forming the Vandermonde matrix\n",
    "nmodes = int(lmax * (lmax+1) / 2)\n",
    "\n",
    "\n",
    "e = paddle.zeros([lmax, mmax, lmax, mmax], dtype=paddle.complex64)\n",
    "for l in range(lmax):\n",
    "    for m in range(l+1):\n",
    "        e[l, m, l, m] = 1.0\n",
    "\n",
    "Sinv = isht(e).to(paddle.complex64) # l x m x nlat x nlon\n",
    "\n",
    "e = paddle.zeros([nlat, nlon, nlat, nlon], dtype=paddle.float32)\n",
    "for lat in range(nlat):\n",
    "    for lon in range(nlon):\n",
    "        e[lat, lon, lat, lon] = 1.0\n",
    "\n",
    "S = sht(e) # nlat x nlon x l x m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = paddle.einsum(\"lmop,lmqr->lopqr\", Sinv, Sinv.conj())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_sphere(out[3, 29, 16].abs(), colorbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS = paddle.einsum(\"lmqr, oplm->qrop\", Sinv, S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(SS.reshape([nlat*nlon, nlat*nlon]).abs(), interpolation=\"nearest\", cmap=\"plasma\")\n",
    "plt.colorbar()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(SS[:,0,:,0].abs(), interpolation=\"nearest\", cmap=\"plasma\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "field = paddle.zeros([nlat, nlon])\n",
    "field[16,3] = 1.0\n",
    "\n",
    "plot_sphere(field, colorbar=True)\n",
    "plot_sphere(isht(sht(field)), colorbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
